import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import numpy as np
import torch
import faiss
from sentence_transformers import SentenceTransformer
from typing import List


class FaissIndexer:
  def __init__(self, embedding_model_path: str = "Qwen/Qwen3-Embedding-0.6B", device: str = "cuda" if torch.cuda.is_available() else "mps"):
    """
    初始化Embedding模型和FAISS索引

    参数:
        embedding_model_path: Embedding模型路径
        device: 计算设备
    """
    self.device = device

    # 加载Embedding模型
    print("⏳ 正在加载Embedding模型...")
    self.embedding_model = SentenceTransformer(
        embedding_model_path,
        device=device
    )

    # 初始化FAISS索引
    self.dim = self.embedding_model.get_sentence_embedding_dimension()
    self.faiss_index = faiss.IndexFlatIP(self.dim)
    self.documents = []

    print("✅ 索引器初始化完成!")

  def index_documents(self, documents: List[str]) -> None:
    """
    为文档集合创建向量索引

    参数:
        documents: 文档列表
    """
    print(f"📚 正在为 {len(documents)} 个文档创建索引...")
    self.documents = documents
    embeddings = self.embedding_model.encode(
        documents,
        convert_to_tensor=True,
        device=self.device
    ).cpu().numpy().astype('float32')
    faiss.normalize_L2(embeddings)
    self.faiss_index.add(embeddings)
    print("✅ 文档索引创建完成!")

  def save_index(self, index_path: str, docs_path: str) -> None:
    """
    保存FAISS索引和文档到磁盘

    参数:
        index_path: FAISS索引保存路径
        docs_path: 文档保存路径
    """
    # 保存FAISS索引
    faiss.write_index(self.faiss_index, index_path)
    print(f"💾 FAISS索引已保存到: {index_path}")

    # 保存文档
    with open(docs_path, 'w', encoding='utf-8') as f:
      for doc in self.documents:
        f.write(doc + '\n')
    print(f"💾 文档已保存到: {docs_path}")

  def update_index(self, new_documents: List[str]) -> None:
    """
    更新向量索引，添加新文档

    参数:
        new_documents: 新的文档列表
    """
    print(f"🔄 正在更新索引，添加 {len(new_documents)} 个新文档...")

    # 编码新文档
    new_embeddings = self.embedding_model.encode(
        new_documents,
        convert_to_tensor=True,
        device=self.device
    ).cpu().numpy().astype('float32')
    faiss.normalize_L2(new_embeddings)

    # 添加到FAISS索引
    self.faiss_index.add(new_embeddings)

    # 更新文档列表
    self.documents.extend(new_documents)

    print(f"✅ 索引更新完成! 现在共有 {len(self.documents)} 个文档")


# ==================== 使用示例 ====================
if __name__ == "__main__":
  # 1. 初始化索引器
  indexer = FaissIndexer(
      embedding_model_path="Qwen/Qwen3-Embedding-0.6B",
      device="cpu"
  )

  # 2. 准备文档库
  documents = [
    "北京是中国的首都，位于华北平原的北部。",
    "巴黎是法国的首都，也是欧洲重要的文化和商业中心。",
    "重力是物体之间相互吸引的力，由艾萨克·牛顿首次系统描述。",
    "量子力学是描述微观粒子行为的物理学理论。",
    "长城是中国古代的军事防御工程，是世界文化遗产。",
    "埃菲尔铁塔位于法国巴黎，是世界上最著名的建筑之一。",
    "黑洞是时空曲率大到连光都无法逃脱的天体。",
    "故宫位于北京，是中国明清两代的皇家宫殿。",
    "相对论是爱因斯坦提出的关于时空和引力的理论。",
    "亚马逊雨林是世界上最大的热带雨林，位于南美洲。",
    "大熊猫是中国的国宝，主要生活在四川的山区。",
    "人工智能是计算机科学的一个分支，旨在创造智能机器。",
    "长城始建于春秋战国时期，用于防御北方游牧民族。",
    "量子纠缠是量子力学中的一种现象，两个粒子无论距离多远都能瞬间影响对方。",
    "卢浮宫位于巴黎，是世界上最大的艺术博物馆之一。"
  ]

  # 3. 创建文档索引
  indexer.index_documents(documents)

  # 4. 保存索引和文档
  indexer.save_index("faiss_index.bin", "documents.txt")